package sample.methodsearch;

import org.apache.commons.lang3.StringUtils;
import org.objectweb.asm.*;
import org.objectweb.asm.util.Printer;

import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

public class MethodFindInvokeDemo {

    public static void main(String[] args) throws IOException {
        invoke(HelloWorld.class, "add", "(II)I");
    }

    private static void invoke(Class<?> clazz, String methodName, String methodDesc) throws IOException {
        ClassReader classReader = new ClassReader(clazz.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES);
        MethodFindInvokeVisitor methodFindInvokeVisitor = new MethodFindInvokeVisitor(Opcodes.ASM9, classWriter, methodName, methodDesc);
        classReader.accept(methodFindInvokeVisitor, ClassReader.EXPAND_FRAMES);
    }

    public static class MethodFindInvokeVisitor extends ClassVisitor {
        private final String methodName;
        private final String methodDesc;

        public MethodFindInvokeVisitor(int api, ClassVisitor classVisitor, String methodName, String methodDesc) {
            super(api, classVisitor);
            this.methodName = methodName;
            this.methodDesc = methodDesc;
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            if (StringUtils.equals(name, methodName) && StringUtils.equals(descriptor, methodDesc)) {
                return new MethodFindInvokeAdapter(api, null);
            }
            return null;
        }

        private static class MethodFindInvokeAdapter extends MethodVisitor {
            private final Set<String> invokeMethods = new HashSet<>();

            public MethodFindInvokeAdapter(int api, MethodVisitor methodVisitor) {
                super(api, methodVisitor);
            }

            @Override
            public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
                String info = String.format("%s %s %s.%s", Printer.OPCODES[opcode], owner, name, descriptor);
                invokeMethods.add(info);
                super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
            }

            @Override
            public void visitEnd() {
                if (!invokeMethods.isEmpty()) {
                    invokeMethods.forEach(System.out::println);
                }
                super.visitEnd();
            }
        }
    }
}
